from uuid import uuid4
from ExtendRegister.db_register import db
from ExtendRegister.model_register import Thesis, Paragraph, GenerateTask
import json
import logging
import shortuuid
from typing import List
from common import Result, ApiResponse, TaskStatus,countCharacters
from common.docx import ThesisDocx
from sqlalchemy import desc
from datetime import datetime

from api.talkgpt.utils import travelOutlineDict
from .schema import (
    ParamSaveSingleParagraph,
    ParamSaveDigest,
    ParamThesisId,
    MovePara,
    SaveNewPara,
    DeleteParagraph,
    SaveThesisProperty,
    ParamGenerateOutline,
)

from common.outline import OutlineKit


LOGGER = logging.getLogger(__name__)


class ThesisServie:

    def deletePara(self, body: DeleteParagraph) -> Result:
        try:
            # 检查
            thesis: Thesis = Thesis.query.filter(Thesis.id == body.thesisId).first()
            if thesis is None:
                return Result().error("未查询到此论文")
            if thesis.uid != body.userId:
                return Result().error("没有查询到此论文")

            para: Paragraph = (
                Paragraph.query.filter(Paragraph.thesisId == body.thesisId)
                .filter(Paragraph.paraId == body.paraId)
                .first()
            )

            if para is None:
                return Result().error(f"没有查询到此段落{body.paraId}")

            # 从tree中删除
            outlineKit = OutlineKit.restoreFromDict(thesis.outline)
            nodeWaitDeleted = outlineKit.findNodeByParaId(body.paraId)
            if nodeWaitDeleted is None:
                return Result().error(f"提纲中没有此段 {body.paraId}")
            allParaIdWaitDeleted = nodeWaitDeleted.getAllChildrenId()
            allParaIdWaitDeleted.append(body.paraId)
            LOGGER.info(f"以下段落将被删除 {allParaIdWaitDeleted}")

            res: Result = outlineKit.deletePara(body.paraId)
            # outlineKit.debugPrint()
            if not res.is_success():
                return res

            # 保存提纲
            thesis.outline = outlineKit.dumpsToDict()
            thesis.save()

            # 删除段落
            Paragraph.query.filter(Paragraph.thesisId == body.thesisId).filter(
                Paragraph.paraId.in_(allParaIdWaitDeleted)
            ).delete()
            db.session.commit()

            return Result()

        except Exception as e:
            LOGGER.exception(e)
        return Result()

    def saveNewPara(self, body: SaveNewPara) -> Result:
        try:
            thesis: Thesis = Thesis.query.filter(Thesis.id == body.thesisId).first()
            if thesis is None:
                return Result().error("未查询到此论文")
            if thesis.uid != body.userId:
                return Result().error("没有查询到此论文")

            outlineKit = OutlineKit.restoreFromDict(thesis.outline)
            if body.addType == "next":
                res: Result = outlineKit.addParaAfter(
                    body.baseParaId, body.newTitle, body.newContent
                )
            else:
                res: Result = outlineKit.addChildPara(
                    body.baseParaId, body.newTitle, body.newContent
                )

            # outlineKit.debugPrint()
            if not res.is_success():
                return res

            newParaId = res.data
        except Exception as e:
            LOGGER.exception(e)
            return Result().error(str(e))

        try:
            Paragraph(
                uid=body.userId,
                thesisId=body.thesisId,
                paraId=newParaId,
                title=body.newTitle,
                paragraph=body.newContent,
            ).save()

            thesis: Thesis = Thesis.query.filter(Thesis.id == body.thesisId).first()
            thesis.outline = outlineKit.dumpsToDict()
            thesis.save()

            return Result()

        except Exception as e:
            LOGGER.exception(e)
            return Result().error(str(e))

    def movelPara(self, body: MovePara) -> Result:
        try:
            thesis: Thesis = Thesis.query.filter(Thesis.id == body.thesisId).first()

            if thesis is None:
                return Result().error("未查询到此论文")

            if thesis.uid != body.userId:
                return ApiResponse().error("没有未查询到此论文").json()

            outlineKit = OutlineKit.restoreFromDict(thesis.outline)

            res: Result = outlineKit.movePara(body.moveType, body.paraId)
            if not res.is_success():
                return res

            thesis: Thesis = Thesis.query.filter(Thesis.id == body.thesisId).first()
            thesis.outline = outlineKit.dumpsToDict()
            thesis.save()

            return Result()
        except Exception as e:
            LOGGER.exception(e)

            return Result().error(str(e))

    def getOutlineProgressById(self, param: ParamGenerateOutline) -> Result:
        thesis: Thesis = Thesis.query.filter(Thesis.id == param.thesisId).first()

        if thesis is None:
            return Result().error("未查询到此论文")

        if thesis.uid != param.userId:
            return ApiResponse().error("未查询到此论文").json()

        outline = thesis.outline

        allParasMap = {}
        paraList: List[Paragraph] = Paragraph.query.filter(
            Paragraph.thesisId == param.thesisId
        ).all()
        for p in paraList:
            allParasMap[p.paraId] = p

        # print(allParasMap.keys())

        wrap = {}
        wrap["totalWait"] = 0

        def _cb(node):
            paraId = node["id"]
            if paraId not in allParasMap:
                # LOGGER.error(f"{paraId} 不在 allParasMap 中")
                node["status"] = "-"
                return

            p: Paragraph = allParasMap[paraId]
            node["status"] = p.status
            node["title"] = p.title

            if p.status == TaskStatus.INIT or p.status == TaskStatus.RUNING:
                wrap["totalWait"] += 1

        outlineWithState = travelOutlineDict(outline, _cb)

        return Result().setData(
            {
                "outline": outlineWithState,
                "totalWait": wrap["totalWait"],
                "thesisStatus": thesis.status,
            }
        )

    def getProgress(self, param: ParamThesisId) -> Result:
        paras: list[Paragraph] = Paragraph.query.filter(
            Paragraph.thesisId == param.thesisId
        ).all()
        result = []
        for para in paras:
            result.append({"title": para.title, "paraId": para.paraId, "state": "s"})
        return Result().setData(result)

    def saveThesisProperty(self, param: SaveThesisProperty) -> Result:
        try:
            thesis: Thesis = Thesis.query.get(param.thesisId)
            if thesis is None:
                return Result().error("thesis not found")

            if thesis.uid != param.userId:
                return Result().error("thesis not found in your list")

            updateDict = {}
            if param.propName == "thanks":
                updateDict["thanks"] = param.content
            elif param.propName == "references":
                updateDict["references"] = param.content
            else:
                return Result().error("propName invalid")

            Thesis.query.filter(Thesis.id == param.thesisId).update(updateDict)
            db.session.commit()
            return Result()
        except Exception as e:
            return Result().error(str(e))

    def saveDigest(self, param: ParamSaveDigest) -> Result:
        try:
            thesis: Thesis = Thesis.query.get(param.thesisId)
            if thesis is None:
                return Result().error("thesis not found")

            if thesis.uid != param.userId:
                return Result().error("thesis not found in your list")

            Thesis.query.filter(Thesis.id == param.thesisId).update(
                {
                    "digest": param.digest,
                    "digestEn": param.digestEn,
                    "keywords": param.keywords,
                    "keywordsEn": param.keywordsEn,
                }
            )
            db.session.commit()
            return Result()
        except Exception as e:
            return Result().error(str(e))

    def saveSingleParagraph(self, param: ParamSaveSingleParagraph) -> Result:
        try:
            thesis: Thesis = Thesis.query.get(param.thesisId)
            if thesis is None:
                return Result().error("thesis not found")

            if thesis.uid != param.userId:
                return Result().error("thesis not found in your list")

            para: Paragraph = (
                Paragraph.query.filter(Paragraph.paraId == param.paragraphId)
                .filter(Paragraph.thesisId == param.thesisId)
                .first()
            )

            if para is None:
                return Result().error(f"paragraph not found {param.paragraphId}")
            Paragraph.query.filter(Paragraph.id == para.id).update(
                {"paragraph": param.text, "title": param.title}
            )
            db.session.commit()
            return Result()
        except Exception as e:
            return Result().error(str(e))

    def exportThesis(self, body: ParamThesisId) -> Result:
        try:

            thesis: Thesis = Thesis.query.get(body.thesisId)
            if thesis is None:
                return Result().error("thesis not found")
            if thesis.uid != body.userId:
                return Result().error("thesis not found in your list")

            res: Result = self.getDetailById(body)
            if not res.is_success():
                return res
            # 上个api返回的outline，中已经带了内容，任然是树形结构
            thesis: Thesis = res.data

            docx = ThesisDocx()
            # 封面
            docx.addCover(thesis.title)
            docx.addPageBreak()
            # 摘要
            docx.addDigest(
                digest=thesis.digest,
                digestEn=thesis.digestEn,
                keywords=thesis.keywords,
                keywordsEn=thesis.keywordsEn,
            )
            docx.addPageBreak()
            # 目录
            docx.addToc()
            docx.addPageBreak()

            def _travel(outline: dict, level: int = 1):
                if "title" in outline and outline["title"] != None:
                    docx.addContentParagraph(outline["title"], outline["text"], level)
                    LOGGER.info(f"add Para:{outline['title']}")
                if "subtitle" in outline:
                    for subOutline in outline["subtitle"]:
                        _travel(subOutline, level + 1)

            # 正文
            _travel(thesis.outline, 0)
            # 参考文件，附录，致谢
            docx.addTail(thesis.references, thesis.thanks)

            now = datetime.strftime(datetime.now(), "%Y%m%d_%H%M%S")
            fileName = f"[神笔][365d.ink]_{thesis.title[:10]}_V{now}.docx"
            docx.export(fileName)

            return Result().setData({"file": fileName})
        except Exception as e:
            LOGGER.exception(e)
            return Result().error(e)

    def deleteThesis(self, body: ParamThesisId) -> Result:
        try:
            thesis: Thesis = Thesis.query.get(body.thesisId)
            if thesis is None:
                return Result().error("thesis not found")
            if thesis.uid != body.userId:
                return Result().error("thesis not found in your list")

            Thesis.query.filter(Thesis.id == body.thesisId).delete()
            Paragraph.query.filter(Paragraph.thesisId == body.thesisId).delete()
            db.session.commit()
            return Result()
        except Exception as e:
            LOGGER.exception(e)
            return Result().error(str(e))

    def getListByUid(self, userid: int) -> Result:
        try:
            thesisList = (
                Thesis.query.filter(Thesis.uid == userid)
                .order_by(desc(Thesis.create_time))
                .all()
            )
            res = []
            for th in thesisList:
                thJson = th.to_json()
                thJson["create_time"] = datetime.strftime(
                    th.create_time, "%Y-%m-%d %H:%M"
                )
                del thJson["outline"]
                res.append(thJson)

            return Result().setData(res)

        except Exception as e:
            LOGGER.exception(e)
            return Result().error(str(e))

    def getThesisById(self, thesisId: int) -> Result:
        try:
            thesis: Thesis = Thesis.query.filter(Thesis.id == thesisId).first()
            if thesis is None:
                return Result().error("未查询到此论文")

            return Result().setData(thesis)
        except Exception as e:
            LOGGER.exception(e)
            return Result().error(f"查询遇到错误{str(e)}")

    def getDetailById(self, body: ParamThesisId) -> Result:
        try:
            thesis: Thesis = Thesis.query.filter(Thesis.id == body.thesisId).first()
            if thesis is None:
                return Result().error("未查询到此论文")

            if thesis.uid != body.userId:
                return Result().error("未查询到此论文2").json()

            allParas = {}
            paragraphs = Paragraph.query.filter(Paragraph.thesisId == thesis.id).all()
            for p in paragraphs:
                allParas[p.paraId] = {"text": p.paragraph, "title": p.title}

            # 把生成的内容放在outline的树结构中
            def _attachParaText(paraDict):
                if paraDict["id"] in allParas:
                    paraDict["text"] = allParas[paraDict["id"]]["text"]
                    paraDict["length"] = countCharacters(paraDict["text"])

                    # 如果用户修改了标题，优先使用，否则使用生成的title
                    user_define_title = allParas[paraDict["id"]]["title"]
                    if user_define_title is not None and user_define_title != "":
                        paraDict["title"] = user_define_title

            thesis.outline = travelOutlineDict(thesis.outline, _attachParaText)
            return Result().setData(thesis)
        except Exception as e:
            LOGGER.exception(e)
            return Result().error(f"查询遇到错误{str(e)}")
